package com.aiwiown.snackmq.cluster.raft;

import com.aiwiown.snackmq.cluster.metadata.TopicMetadata;
import com.aiwiown.snackmq.cluster.routing.Binding;
import com.aiwiown.snackmq.cluster.routing.Exchange;
import com.aiwiown.snackmq.cluster.routing.exchange.DirectExchange;
import com.aiwiown.snackmq.common.protocol.MetadataOperation;
import com.aiwiown.snackmq.common.serialization.JsonSerializer;
import com.aiwiown.snackmq.common.serialization.SerializerFactory;
import com.alipay.sofa.jraft.Closure;
import com.alipay.sofa.jraft.Iterator;
import com.alipay.sofa.jraft.Status;
import com.alipay.sofa.jraft.core.StateMachineAdapter;
import com.alipay.sofa.jraft.error.RaftError;
import com.alipay.sofa.jraft.storage.snapshot.SnapshotReader;
import com.alipay.sofa.jraft.storage.snapshot.SnapshotWriter;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.IOException;
import java.io.Serializable;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Collections;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * Raft 协议的状态机实现类，用于管理元数据。
 * 该类继承自 StateMachineAdapter，并实现了其抽象方法。
 */
@Slf4j
public class MetadataStateMachine extends StateMachineAdapter {

    // 定义元数据快照文件名
    private static final String METADATA_SNAPSHOT_FILENAME = "metadata.json";
    // 使用 JSON 序列化器
    private final JsonSerializer serializer = SerializerFactory.getJsonSerializer();

    // 交换机和主题的并发映射，用于存储元数据
    private final ConcurrentMap<String, Exchange> exchanges = new ConcurrentHashMap<>();
    private final ConcurrentMap<String, TopicMetadata> topics = new ConcurrentHashMap<>();

    /**
     * 处理日志条目应用到状态机。
     * 该方法重写了父类的 onApply 方法，用于处理元数据操作。
     *
     * @param iter 日志条目迭代器
     */
    @Override
    public void onApply(Iterator iter) {
        while (iter.hasNext()) {
            Closure done = iter.done();
            try {
                MetadataOperation op = serializer.deserialize(iter.getData().array(), MetadataOperation.class);
                log.info("Applying metadata operation: {}", op);

                switch (op.getOperationType()) {
                    case CREATE_TOPIC:
                        handleCreateTopic(op);
                        break;
                    case DECLARE_EXCHANGE:
                        handleDeclareExchange(op);
                        break;
                    case BIND_QUEUE:
                        handleBindQueue(op);
                        break;
                    default:
                        log.warn("Unknown metadata operation type: {}", op.getOperationType());
                }

                if (done != null) {
                    done.run(Status.OK());
                }
            } catch (Exception e) {
                log.error("Failed to apply metadata operation at index {}", iter.getIndex(), e);
                if (done != null) {
                    done.run(new Status(RaftError.EIO, "Failed to apply metadata operation: " + e.getMessage()));
                }
            } finally {
                iter.next();
            }
        }
    }

    /**
     * 处理创建主题的操作。
     * 如果主题已存在，则忽略该操作。
     * 否则，创建一个新的 TopicMetadata 实例，并为其分配空的分区信息。
     *
     * @param op 元数据操作
     */
    private void handleCreateTopic(MetadataOperation op) {
        if (topics.containsKey(op.getTopicName())) {
            log.warn("Topic '{}' already exists. Ignoring CREATE_TOPIC operation.", op.getTopicName());
            return;
        }
        // 使用新的单参数构造函数创建 TopicMetadata
        TopicMetadata topicMetadata = new TopicMetadata(op.getTopicName());

        // 为每个分区创建空的分配方案
        for (int i = 0; i < op.getPartitions(); i++) {
            TopicMetadata.PartitionAssignment emptyAssignment = new TopicMetadata.PartitionAssignment(Collections.emptyList(), null);
            topicMetadata.setPartitionAssignment(i, emptyAssignment);
        }

        topics.put(op.getTopicName(), topicMetadata);
        log.info("State Machine: Created new topic '{}' with {} partitions.", op.getTopicName(), op.getPartitions());
    }

    /**
     * 处理声明交换机的操作。
     * 如果交换机已存在，则忽略该操作。
     * 目前只支持 DirectExchange 类型的交换机。
     *
     * @param op 元数据操作
     */
    private void handleDeclareExchange(MetadataOperation op) {
        if (exchanges.containsKey(op.getExchangeName())) {
            log.warn("Exchange '{}' already exists. Ignoring DECLARE_EXCHANGE operation.", op.getExchangeName());
            return;
        }
        Exchange exchange = new DirectExchange(op.getExchangeName());
        exchanges.put(op.getExchangeName(), exchange);
        log.info("State Machine: Creating new exchange '{}' of type {}", op.getExchangeName(), op.getExchangeType());
    }

    /**
     * 处理绑定队列的操作。
     * 如果交换机不存在，则记录错误并返回。
     * 否则，创建一个新的 Binding 实例并将其添加到交换机中。
     *
     * @param op 元数据操作
     */
    private void handleBindQueue(MetadataOperation op) {
        Exchange exchange = exchanges.get(op.getExchangeName());
        if (exchange == null) {
            log.error("Cannot bind queue '{}' because exchange '{}' does not exist.", op.getQueueName(), op.getExchangeName());
            return;
        }
        Binding binding = new Binding(op.getExchangeName(), op.getQueueName(), op.getRoutingKey());
        exchange.addBinding(binding);
        log.info("State Machine: Bound queue '{}' to exchange '{}' with routing key '{}'", op.getQueueName(), op.getExchangeName(), op.getRoutingKey());
    }

    /**
     * 获取交换机。
     *
     * @param name 交换机名称
     * @return 交换机实例，如果不存在则返回 null
     */
    public Exchange getExchange(String name) {
        return this.exchanges.get(name);
    }

    /**
     * 获取主题。
     *
     * @param name 主题名称
     * @return 主题元数据实例，如果不存在则返回 null
     */
    public TopicMetadata getTopic(String name) {
        return this.topics.get(name);
    }

    /**
     * 保存快照。
     * 将当前的交换机和主题元数据保存到指定的快照写入器中。
     *
     * @param writer 快照写入器
     * @param done   完成回调
     */
    @Override
    public void onSnapshotSave(SnapshotWriter writer, Closure done) {
        log.info("Saving metadata snapshot to {}", writer.getPath());
        try {
            MetadataSnapshot snapshot = new MetadataSnapshot(this.exchanges, this.topics);
            byte[] snapshotBytes = serializer.serialize(snapshot);
            String snapshotPath = writer.getPath() + File.separator + METADATA_SNAPSHOT_FILENAME;
            Files.write(Paths.get(snapshotPath), snapshotBytes);
            writer.addFile(METADATA_SNAPSHOT_FILENAME);
            done.run(Status.OK());
            log.info("Metadata snapshot saved successfully.");
        } catch (Exception e) {
            log.error("Failed to save metadata snapshot", e);
            done.run(new Status(RaftError.EIO, "Failed to save metadata snapshot: " + e.getMessage()));
        }
    }

    /**
     * 加载快照。
     * 从指定的快照读取器中加载交换机和主题元数据。
     *
     * @param reader 快照读取器
     * @return true 表示加载成功，false 表示加载失败
     */
    @Override
    public boolean onSnapshotLoad(SnapshotReader reader) {
        log.info("Loading metadata snapshot from {}", reader.getPath());
        if (reader.getFileMeta(METADATA_SNAPSHOT_FILENAME) == null) {
            log.warn("Metadata snapshot file '{}' not found, starting with a clean state.", METADATA_SNAPSHOT_FILENAME);
            return true;
        }
        try {
            String snapshotPath = reader.getPath() + File.separator + METADATA_SNAPSHOT_FILENAME;
            byte[] snapshotBytes = Files.readAllBytes(Paths.get(snapshotPath));
            MetadataSnapshot snapshot = serializer.deserialize(snapshotBytes, MetadataSnapshot.class);

            if (snapshot != null) {
                this.exchanges.clear();
                this.topics.clear();
                if (snapshot.getExchanges() != null) {
                    this.exchanges.putAll(snapshot.getExchanges());
                }
                if (snapshot.getTopics() != null) {
                    this.topics.putAll(snapshot.getTopics());
                }
                log.info("Metadata snapshot loaded successfully. Restored {} exchanges and {} topics.", this.exchanges.size(), this.topics.size());
            }
            return true;
        } catch (IOException e) {
            log.error("Failed to load metadata snapshot.", e);
            return false;
        }
    }

    /**
     * 元数据快照类，用于序列化和反序列化元数据。
     */
    @Getter
    @NoArgsConstructor
    @AllArgsConstructor
    public static class MetadataSnapshot implements Serializable {
        private static final long serialVersionUID = -8456110983272801486L;
        private ConcurrentMap<String, Exchange> exchanges;
        private ConcurrentMap<String, TopicMetadata> topics;
    }
}
